# Standard libraries
import os
import time
import warnings
import re

# Data handling
import numpy as np
import pandas as pd
from PIL import Image

# Visualization
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib

# PyTorch basics
import torch
import torch.nn as nn
import torch.optim as optim

# PyTorch helpers for data loading and transforms
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
import torchvision.utils as vutils

# Metrics and evaluation
from torchmetrics.image.fid import FrechetInceptionDistance
from torchmetrics.image import InceptionScore
from torchsummary import summary

# Progress bar
from tqdm import tqdm

# Ignore warnings
warnings.filterwarnings('ignore')

# Change theme of charts
sns.set_theme(style='darkgrid')
# Change font of charts
sns.set(font='Century Gothic')
# Variable for color palettes
color_palette = sns.color_palette('muted')
# Increase embed limit for animation
matplotlib.rcParams['animation.embed_limit'] = 100

class CatImagesDataset(Dataset):
    def __init__(self, directory, transform=None, limit=2000):
        self.directory = directory
        self.transform = transform
        self.image_paths = [os.path.join(directory, fname) for fname in os.listdir(directory) if fname.endswith('.jpg')][:limit]

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        with open(img_path, 'rb') as f:
            img = Image.open(f)
            img = img.convert('RGB')
            if self.transform:
                img = self.transform(img)
        return img
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
cat_dataset = CatImagesDataset('./cat_ca2/cats', transform=transform)

batch_size = 64
cat_dataloader = DataLoader(cat_dataset, batch_size=batch_size, shuffle=True)

images = []
for batch in cat_dataloader:
    images.append(batch)

images = torch.cat(images)

IMAGE_SIZE = images[0].size()[1]
CRITIC_ITERATIONS = 5
WEIGHT_CLIP = 0.01
LR = 2e-4
EPOCHS = 200
nz = 100
ndf = 128
ngf = 128
nc = 3
ngpu = 1
real_label = 1
fake_label = 0
color_palette = sns.color_palette('muted')
device = torch.device('cuda' if (torch.cuda.is_available() and ngpu > 0) else 'cpu')
print('Device used:', device)

Device used: cuda

def normalize_images(images):
    normalized_images = [(image + 1) / 2 for image in images]
    return normalized_images

def plot_images(images, rows=2, cols=8):
    fig, axs = plt.subplots(rows, cols, figsize=(cols*3, rows*3))
    axs = axs.ravel()
    for i in range(min(len(images), rows * cols)):
        axs[i].imshow(images[i].permute(1, 2, 0)) 
        axs[i].axis('off')
    plt.show()

plot_images(normalize_images(images))

# Monitor Class for GANs
class GANMonitor:
    def __init__(self, patience=5):
        self.G_losses = []
        self.D_losses = []
        self.fid_scores = []
        self.is_scores = []
        self.best_fid_iter = 0
        self.kl_divergence = None
        self.best_fid = float('inf')
        self.improve = True
 
    def plot_loss(self):
        # Plotting
        plt.figure(figsize=(10,6))
        sns.lineplot(self.G_losses, label='Generator Loss', color=color_palette[2])
        sns.lineplot(self.D_losses, label='Discriminator Loss', color=color_palette[3])
        plt.title('Generator vs Discriminator Loss')
        plt.xlabel('Iterations')
        plt.ylabel('Loss')
        plt.legend()
        plt.show()
    
    @staticmethod
    def transform(tensor):
        tensor = (tensor + 1) / 2  # Rescale to [0, 1]
        return tensor

    def calc_fid(self,real,fake,i):
        fid = FrechetInceptionDistance(feature=2048, normalize=True, reset_real_features=True).to('cuda')
        real_images_norm = self.transform(real.to('cuda'))
        fid.update(real_images_norm, real=True)
        fake_images_norm = self.transform(fake.to('cuda'))
        fid.update(fake_images_norm, real=False)
        fid_score = fid.compute()
        if fid_score < self.best_fid:
            self.best_fid = fid_score
            self.best_fid_iter = i
            self.improve = True
        else:
            self.improve = False
        return fid_score
    
    def calc_is(self,imgs):
        inception = InceptionScore(normalize=True)
        norm_imgs = self.transform(imgs.to('cpu'))
        scaled_imgs = (norm_imgs * 255).clamp(0, 255).to(torch.uint8)
        inception.update(scaled_imgs)
        return inception.compute()
    
    def plot_scores(self):
        # Convert tensors to lists using list comprehension
        fid_scores = [score.item() for score in self.fid_scores]
        inception_scores, standard_deviations = zip(*[(score.item(), std.item()) for score, std in self.is_scores])

        # Print best results
        min_fid = min(fid_scores)
        max_is = max(inception_scores)
        
        # Print results
        print(f"Minimum FID Score of {min_fid} obtained at iteration of {self.best_fid_iter}.")
        print(f"Maximum IS Score of {max_is}.")

        # Create a figure with three subplots
        fig, axes = plt.subplots(1, 3, figsize=(16, 5))

        # Define common properties for line plots
        plot_props = {'marker': 'o', 'linestyle': '-', 'x': range(len(fid_scores))}

        # FID Scores Line Plot
        sns.lineplot(y=fid_scores, color=color_palette[0], ax=axes[0], **plot_props)
        axes[0].set(title="FID Scores", xlabel="Index", ylabel="FID Score")

        # Inception Scores Line Plot
        sns.lineplot(y=inception_scores, label='Inception Score', color=color_palette[1], ax=axes[1], **plot_props)
        axes[1].set(title='Inception Scores', xlabel='Index', ylabel='Inception Score')
        axes[1].legend()
        axes[1].grid(True)

        # Standard Deviations Line Plot
        sns.lineplot(y=standard_deviations, label='Standard Deviation', color=color_palette[2], ax=axes[2], **plot_props)
        axes[2].set(title='Inception Standard Deviations', xlabel='Index', ylabel='Standard Deviation')
        axes[2].legend()
        axes[2].grid(True)

        # Change suptitle
        fig.suptitle(t='Fréchet Inception Distance and Inception Scores')

        # Show the plot
        plt.tight_layout()
        plt.show()

    def save_weights(self,i,netG, critic, filepath):
        if self.improve:
            # Define your desired file path
            model_dir = filepath
            # Check if the directory exists, and create it if it doesn't
            if not os.path.exists(model_dir):
                os.makedirs(model_dir)
            # Save generator and discriminator weights
            torch.save(netG, filepath + f"/generator-{i}.pth")
            torch.save(critic, filepath + f"/critic-{i}.pth")

    def store_print_metrics(self, real, fake,i):
        current_fid_score = self.calc_fid(real=real, fake=fake, i=i)
        current_is_score = self.calc_is(imgs=fake)
        self.fid_scores.append(current_fid_score)
        self.is_scores.append(current_is_score)
        print(f"Current scores at iteration {i} | FID: {current_fid_score} | IS: {current_is_score[0].item()}")


# ============
# |   WGAN   |
# ============

class WGANGenerator(nn.Module):
    def __init__(self):
        super(WGANGenerator, self).__init__()

        self.model = nn.Sequential(
            self._block(nz, ngf * 8, 4, 1, 0),
            self._block(ngf * 8, ngf * 4), 
            self._block(ngf * 4, ngf * 2), 
            self._block(ngf * 2, ngf),
            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def _block(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, batch_norm=True):
        layers = [nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=not batch_norm)]
        if batch_norm:
            layers.append(nn.BatchNorm2d(out_channels))
        layers.append(nn.LeakyReLU(0.2))

        return nn.Sequential(*layers)
        
    def forward(self, z):
        return self.model(z)
    
# WGAN Critic
class WGANCritic(nn.Module):
    def __init__(self):
        super(WGANCritic, self).__init__()

        self.model = nn.Sequential(
            self._block(nc, ndf, batch_norm=False),
            self._block(ndf, ndf * 2),
            self._block(ndf * 2, ndf * 4),
            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def _block(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, batch_norm=True):
        layers = [nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=not batch_norm)]
        if batch_norm:
            layers.append(nn.BatchNorm2d(out_channels, momentum=0.9))
        layers.append(nn.ReLU(inplace=True))
        return nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

class WGANTrainer(GANMonitor):
    def __init__(self, critic, netG, dataloader, device, fixed_noise, real_label, fake_label, nz, optimizerC, optimizerG, num_epochs, weight_clip, critic_iterations):
        super().__init__()
        self.critic = critic
        self.netG = netG
        self.dataloader = dataloader
        self.device = device
        self.fixed_noise = fixed_noise
        self.real_label = real_label
        self.fake_label = fake_label
        self.nz = nz
        self.optimizerC = optimizerC
        self.optimizerG = optimizerG
        self.num_epochs = num_epochs
        self.weight_clip = weight_clip
        self.critic_iterations = critic_iterations
        self.file_path = "./wgan_weights"
        self.temp_errC = []
        self.iters = 0
        self.img_list = []

    def train(self):
        print("Starting Training Loop...")
        for epoch in tqdm(range(self.num_epochs), desc="Training Progress"):
            for i, data in enumerate(self.dataloader, 0):
                start_time = time.time()
                # Train discriminator
                errC, errG, real, fake = self.train_step(data)
                # Train generator
                end_time = time.time()
                iteration_time = end_time - start_time
                # At every 50 iterations, record current loss of the Generator
                if i % self.critic_iterations == 0:
                    print('[%d/%d][%d/%d] \tLoss_D: %.4f | Loss_G: %.4f | Iteration Time: %.4f sec'
                % (epoch, self.num_epochs, i, len(self.dataloader),
                    errC.item(), errG.item(), iteration_time))
                self.check_generator_progress(epoch, i, real, fake)

    def train_step(self, data):
        # Train critic for n critic iterations
        for _ in range(self.critic_iterations):
            # Reset gradients for Discriminator
            self.critic.zero_grad()
            # Real images
            real = data.to(self.device)
            # Noise
            noise = torch.randn(data.size(0), self.nz, 1, 1, device=self.device)
            # Fake image batch generated with G
            fake_data = self.netG(noise)
            # Forward pass real batch through C
            real_output = self.critic(real).view(-1)
            # Output of C when given fake/generated data
            fake_output = self.critic(fake_data.detach()).view(-1)
            # Calculate loss for C
            errC = -(torch.mean(real_output) - torch.mean(fake_output))
            # Append to temporary array
            self.temp_errC.append(errC)
            # Backpropagate for C
            errC.backward()
            self.optimizerC.step()
            # Weight clipping for C
            for p in self.critic.parameters():
                p.data.clamp_(-self.weight_clip, self.weight_clip)
            
        # Reset gradients for Generator
        self.netG.zero_grad()
        # Since we just updated D, perform another forward pass of fake_data through D
        fake_output = self.critic(fake_data).view(-1)
        # Calculate error for G
        errG = -torch.mean(fake_output)
        # Store generator loss
        self.G_losses.append(errG.item())
        # Store average critic loss
        self.D_losses.append((sum(self.temp_errC)/len(self.temp_errC)).item())
        # Reset temp errC array
        self.temp_errC = []
        # Backpropagate for G
        errG.backward()
        self.optimizerG.step()
        return errC, errG, real, fake_data
    
    def check_generator_progress(self, epoch, i, real, fake):
        if (self.iters % 500 == 0) or ((epoch == self.num_epochs-1) and (i == len(self.dataloader)-1)):
            with torch.no_grad():
                # Calculate, store and print FID and IS score
                self.store_print_metrics(real=real, fake=fake, i=self.iters)
                # Save weights
                self.save_weights(self.iters,self.netG, self.critic, self.file_path)
                # Generate image and create image grid
                fake = self.netG(self.fixed_noise).detach().cpu()
                generated_image = np.transpose(vutils.make_grid(fake, padding=2, normalize=True), (1,2,0))
                # Plot image grid
                self.img_list.append(generated_image)
                plt.figure(figsize=(20,6))
                plt.axis("off")
                plt.title(f"Generated Images at {self.iters} iterations")
                plt.imshow(generated_image)
                plt.show()
        self.iters += 1

# ==============
# |   WGAN-GP  |
# ==============
    
class WGANGPGenerator(nn.Module):
    def __init__(self, nz=100, ngf=64, nc=3):
        super(WGANGPGenerator, self).__init__()

        self.model = nn.Sequential(
            self._block(nz, ngf * 8, 4, 1, 0),
            self._block(ngf * 8, ngf * 4), 
            self._block(ngf * 4, ngf * 2), 
            self._block(ngf * 2, ngf),
            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def _block(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, batch_norm=True):
        layers = [nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=not batch_norm)]
        if batch_norm:
            layers.append(nn.BatchNorm2d(out_channels))
        layers.append(nn.LeakyReLU(0.2))

        return nn.Sequential(*layers)
    
    def forward(self, z):
        return self.model(z)

    
# WGAN Critic
class WGANGPCritic(nn.Module):
    def __init__(self):
        super(WGANGPCritic, self).__init__()

        self.model = nn.Sequential(
            self._block(nc, ndf, batch_norm=False),
            self._block(ndf, ndf * 2),
            self._block(ndf * 2, ndf * 4),
            nn.Conv2d(ndf * 4, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def _block(self, in_channels, out_channels, kernel_size=4, stride=2, padding=1, batch_norm=True):
        layers = [nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=not batch_norm)]
        if batch_norm:
            layers.append(nn.InstanceNorm2d(out_channels, momentum=0.9))
        layers.append(nn.ReLU(inplace=True))
        return nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)
    
class WGANGPTrainer(GANMonitor):
    def __init__(self, critic, netG, dataloader, device, fixed_noise, real_label, fake_label, nz, optimizerC, optimizerG, num_epochs, critic_iterations, lambda_gp):
        super().__init__()
        self.critic = critic
        self.netG = netG
        self.dataloader = dataloader
        self.device = device
        self.fixed_noise = fixed_noise
        self.real_label = real_label
        self.fake_label = fake_label
        self.nz = nz
        self.optimizerC = optimizerC
        self.optimizerG = optimizerG
        self.num_epochs = num_epochs
        self.critic_iterations = critic_iterations
        self.lambda_gp = lambda_gp
        self.file_path = "./wgan-gp_weights"
        self.temp_errC = []
        self.iters = 0
        self.img_list = []

    def train(self):
        print("Starting Training Loop...")
        for epoch in tqdm(range(self.num_epochs), desc="Training Progress"):
            for i, data in enumerate(self.dataloader, 0):
                start_time = time.time()
                # Train discriminator
                errC, errG, real, fake = self.train_step(data)
                # Train generator
                end_time = time.time()
                iteration_time = end_time - start_time
                # At every 50 iterations, record current loss of the Generator
                if i % self.critic_iterations == 0:
                    print('[%d/%d][%d/%d] \tLoss_D: %.4f | Loss_G: %.4f | Iteration Time: %.4f sec'
                % (epoch, self.num_epochs, i, len(self.dataloader),
                    errC.item(), errG.item(), iteration_time))
                self.check_generator_progress(epoch, i, real, fake)

    def train_step(self, data):
        # Train critic for n critic iterations
        for _ in range(self.critic_iterations):
            # Reset gradients for Discriminator
            self.critic.zero_grad()
            # Real images
            real = data.to(self.device)
            # Noise
            noise = torch.randn(data.size(0), self.nz, 1, 1, device=self.device)
            # Fake image batch generated with G
            fake_data = self.netG(noise)
            # Forward pass real batch through C
            real_output = self.critic(real).view(-1)
            # Output of C when given fake/generated data
            fake_output = self.critic(fake_data.detach()).view(-1)
            # Gradient penalty
            gp = self.gradient_penalty(self.critic, real, fake_data, device=self.device)
            # Calculate loss for C
            errC = (-(torch.mean(real_output) - torch.mean(fake_output)) + self.lambda_gp * gp)
            # Append to temporary array
            self.temp_errC.append(errC)
            # Backpropagate for C
            errC.backward(retain_graph=True)
            self.optimizerC.step()
    
        # Reset gradients for Generator
        self.netG.zero_grad()
        # Since we just updated D, perform another forward pass of fake_data through D
        fake_output = self.critic(fake_data).view(-1)
        # Calculate error for G
        errG = -torch.mean(fake_output)
        # Store generator loss
        self.G_losses.append(errG.item())
        # Store average critic loss
        self.D_losses.append((sum(self.temp_errC)/len(self.temp_errC)).item())
        # Reset temp errC array
        self.temp_errC = []
        # Backpropagate for G
        errG.backward()
        self.optimizerG.step()

        return errC, errG, real, fake_data
    
    def check_generator_progress(self, epoch, i, real, fake):
        if (self.iters % 500 == 0) or ((epoch == self.num_epochs-1) and (i == len(self.dataloader)-1)):
            with torch.no_grad():
                # Calculate, store and print FID and IS score
                self.store_print_metrics(real=real, fake=fake, i=self.iters)
                # Save weights
                self.save_weights(self.iters,self.netG, self.critic, self.file_path)
                # Generate image and create image grid
                fake = self.netG(self.fixed_noise).detach().cpu()
                generated_image = np.transpose(vutils.make_grid(fake, padding=2, normalize=True), (1,2,0))
                # Plot image grid
                self.img_list.append(generated_image)
                plt.figure(figsize=(20,6))
                plt.axis("off")
                plt.title(f"Generated Images at {self.iters} iterations")
                plt.imshow(generated_image)
                plt.show()
        self.iters += 1

    @staticmethod
    def gradient_penalty(critic, real, fake, device='cpu'):
        batch_size, C, H, W = real.shape
        epsilon = torch.rand((batch_size, 1, 1, 1)).repeat(1, C, H, W).to(device) # Create interpolated images
        interpolated_images = real * epsilon + fake * (1 - epsilon) # Interpolate real image with fake image

        # Calculate critic score
        mixed_scores = critic(interpolated_images)
        gradient = torch.autograd.grad(
            inputs=interpolated_images,
            outputs=mixed_scores,
            grad_outputs=torch.ones_like(mixed_scores),
            create_graph=True,
            retain_graph=True
        )[0]
        gradient = gradient.view(gradient.shape[0], -1)
        gradient_norm = gradient.norm(2, dim=1)
        gradient_penalty = torch.mean((gradient_norm-1)**2)

        return gradient_penalty

# Create generator
netG = WGANGenerator().to(device=device)
# Print generator
print(netG)

WGANGenerator(
  (model): Sequential(
    (0): Sequential(
      (0): ConvTranspose2d(100, 1024, kernel_size=(4, 4), stride=(1, 1), bias=False)
      (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (1): Sequential(
      (0): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (2): Sequential(
      (0): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (3): Sequential(
      (0): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (4): ConvTranspose2d(128, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (5): Tanh()
  )
)

# Create Critic
critic = WGANCritic().to(device=device)
# Print critic
print(critic)

WGANCritic(
  (model): Sequential(
    (0): Sequential(
      (0): Conv2d(3, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
      (1): ReLU(inplace=True)
    )
    (1): Sequential(
      (0): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
    (2): Sequential(
      (0): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.9, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
    )
    (3): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1), bias=False)
    (4): Sigmoid()
  )
)

# Optimizers
optimizerC = optim.RMSprop(critic.parameters(), lr=LR, alpha=0.9)
optimizerG = optim.RMSprop(netG.parameters(), lr=LR, alpha=0.9)
# Fixed noise (latent vectors)
fixed_noise = torch.randn(IMAGE_SIZE, nz, 1, 1, device=device)

wgan_trainer = WGANTrainer(critic, netG, cat_dataloader, device, fixed_noise, real_label, fake_label, nz, optimizerC, optimizerG, EPOCHS, WEIGHT_CLIP, CRITIC_ITERATIONS)
wgan_trainer.train()

Starting Training Loop...

Training Progress:   0%|          | 0/200 [00:00<?, ?it/s]

[0/200][0/32] 	Loss_D: -0.0061 | Loss_G: -0.4950 | Iteration Time: 0.8203 sec
Current scores at iteration 0 | FID: 399.6596374511719 | IS: 1.0986696481704712

[0/200][5/32] 	Loss_D: -0.0355 | Loss_G: -0.4774 | Iteration Time: 0.0926 sec
[0/200][10/32] 	Loss_D: -0.0771 | Loss_G: -0.4552 | Iteration Time: 0.0915 sec
[0/200][15/32] 	Loss_D: -0.1251 | Loss_G: -0.4448 | Iteration Time: 0.0910 sec
[0/200][20/32] 	Loss_D: -0.1483 | Loss_G: -0.4428 | Iteration Time: 0.0895 sec
[0/200][25/32] 	Loss_D: -0.1391 | Loss_G: -0.4183 | Iteration Time: 0.0880 sec

Training Progress:   0%|          | 1/200 [00:10<34:12, 10.32s/it]

[0/200][30/32] 	Loss_D: -0.1369 | Loss_G: -0.4088 | Iteration Time: 0.0905 sec
[1/200][0/32] 	Loss_D: -0.1676 | Loss_G: -0.4215 | Iteration Time: 0.0890 sec
[1/200][5/32] 	Loss_D: -0.1389 | Loss_G: -0.4075 | Iteration Time: 0.0890 sec
[1/200][10/32] 	Loss_D: -0.1614 | Loss_G: -0.4085 | Iteration Time: 0.0890 sec
[1/200][15/32] 	Loss_D: -0.1779 | Loss_G: -0.4080 | Iteration Time: 0.0895 sec
[1/200][20/32] 	Loss_D: -0.1684 | Loss_G: -0.4107 | Iteration Time: 0.0905 sec
[1/200][25/32] 	Loss_D: -0.1813 | Loss_G: -0.4132 | Iteration Time: 0.0900 sec

Training Progress:   1%|          | 2/200 [00:14<21:25,  6.49s/it]

[1/200][30/32] 	Loss_D: -0.1916 | Loss_G: -0.4038 | Iteration Time: 0.0890 sec
[2/200][0/32] 	Loss_D: -0.1625 | Loss_G: -0.4080 | Iteration Time: 0.0895 sec
[2/200][5/32] 	Loss_D: -0.1826 | Loss_G: -0.4075 | Iteration Time: 0.0885 sec
[2/200][10/32] 	Loss_D: -0.1810 | Loss_G: -0.4110 | Iteration Time: 0.0915 sec
[2/200][15/32] 	Loss_D: -0.0916 | Loss_G: -0.4696 | Iteration Time: 0.0930 sec
[2/200][20/32] 	Loss_D: -0.1654 | Loss_G: -0.4080 | Iteration Time: 0.0900 sec
[2/200][25/32] 	Loss_D: -0.1647 | Loss_G: -0.4103 | Iteration Time: 0.0895 sec

Training Progress:   2%|▏         | 3/200 [00:17<17:19,  5.28s/it]

[2/200][30/32] 	Loss_D: -0.1698 | Loss_G: -0.4166 | Iteration Time: 0.0857 sec
[3/200][0/32] 	Loss_D: -0.1822 | Loss_G: -0.4110 | Iteration Time: 0.0900 sec
[3/200][5/32] 	Loss_D: -0.1496 | Loss_G: -0.4056 | Iteration Time: 0.0905 sec
[3/200][10/32] 	Loss_D: -0.1800 | Loss_G: -0.4053 | Iteration Time: 0.0895 sec
[3/200][15/32] 	Loss_D: -0.1468 | Loss_G: -0.4034 | Iteration Time: 0.0895 sec
[3/200][20/32] 	Loss_D: -0.1724 | Loss_G: -0.4093 | Iteration Time: 0.0890 sec
[3/200][25/32] 	Loss_D: -0.1794 | Loss_G: -0.4135 | Iteration Time: 0.0910 sec

wgan_trainer.plot_loss()

wgan_trainer.plot_scores()

Minimum FID Score of 160.5990753173828 obtained at iteration of 6000.
Maximum IS Score of 2.464463710784912.

# Create generator
netG = WGANGPGenerator().to(device=device)
# Print generator
print(netG)

WGANGPGenerator(
  (model): Sequential(
    (0): Sequential(
      (0): ConvTranspose2d(100, 512, kernel_size=(4, 4), stride=(1, 1), bias=False)
      (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (1): Sequential(
      (0): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (2): Sequential(
      (0): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (3): Sequential(
      (0): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): LeakyReLU(negative_slope=0.2)
    )
    (4): ConvTranspose2d(64, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    (5): Tanh()
  )
)

# Create Critic
critic = WGANGPCritic().to(device=device)
# Print critic
print(critic)

WGANGPCritic(
  (model): Sequential(
    (0): Sequential(
      (0): Conv2d(3, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
      (1): ReLU(inplace=True)
    )
    (1): Sequential(
      (0): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): InstanceNorm2d(256, eps=1e-05, momentum=0.9, affine=False, track_running_stats=False)
      (2): ReLU(inplace=True)
    )
    (2): Sequential(
      (0): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
      (1): InstanceNorm2d(512, eps=1e-05, momentum=0.9, affine=False, track_running_stats=False)
      (2): ReLU(inplace=True)
    )
    (3): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1), bias=False)
    (4): Sigmoid()
  )
)

# Optimizers
optimizerC = optim.Adam(critic.parameters(), lr=LR, betas=(0.0, 0.9)) 
optimizerG = optim.Adam(netG.parameters(), lr=LR, betas=(0.0, 0.9))
# Fixed noise (latent vectors)
fixed_noise = torch.randn(IMAGE_SIZE, nz, 1, 1, device=device)
# Lambda
lambda_gp = 10

wgangp_trainer = WGANGPTrainer(critic, netG, cat_dataloader, device, fixed_noise, real_label, fake_label, nz, optimizerC, optimizerG, EPOCHS, CRITIC_ITERATIONS, lambda_gp)
wgangp_trainer.train()

Starting Training Loop...

Training Progress:   0%|          | 0/200 [00:00<?, ?it/s]

[0/200][0/32] 	Loss_D: 0.3098 | Loss_G: -0.4351 | Iteration Time: 1.6757 sec
Current scores at iteration 0 | FID: 363.68902587890625 | IS: 1.1739404201507568

[0/200][5/32] 	Loss_D: 0.0249 | Loss_G: -0.2451 | Iteration Time: 1.1456 sec
[0/200][10/32] 	Loss_D: -0.2798 | Loss_G: -0.2202 | Iteration Time: 0.9056 sec
[0/200][15/32] 	Loss_D: 0.3199 | Loss_G: -0.3168 | Iteration Time: 1.1264 sec
[0/200][20/32] 	Loss_D: -0.2145 | Loss_G: -0.2528 | Iteration Time: 0.9751 sec
[0/200][25/32] 	Loss_D: -0.2108 | Loss_G: -0.3601 | Iteration Time: 1.2558 sec

Training Progress:   0%|          | 1/200 [00:43<2:22:58, 43.11s/it]

[0/200][30/32] 	Loss_D: -0.1920 | Loss_G: -0.3887 | Iteration Time: 0.9213 sec
[1/200][0/32] 	Loss_D: -0.0512 | Loss_G: -0.3675 | Iteration Time: 0.9041 sec
[1/200][5/32] 	Loss_D: 0.0098 | Loss_G: -0.3825 | Iteration Time: 1.0291 sec
[1/200][10/32] 	Loss_D: -0.1522 | Loss_G: -0.3452 | Iteration Time: 1.2386 sec
[1/200][15/32] 	Loss_D: -0.2318 | Loss_G: -0.3863 | Iteration Time: 0.9996 sec
[1/200][20/32] 	Loss_D: -0.0175 | Loss_G: -0.3134 | Iteration Time: 0.9616 sec
[1/200][25/32] 	Loss_D: -0.1546 | Loss_G: -0.3689 | Iteration Time: 0.9581 sec

Training Progress:   1%|          | 2/200 [01:16<2:04:04, 37.60s/it]

[1/200][30/32] 	Loss_D: -0.1740 | Loss_G: -0.2294 | Iteration Time: 1.3574 sec
[2/200][0/32] 	Loss_D: -0.2770 | Loss_G: -0.3015 | Iteration Time: 0.9561 sec
[2/200][5/32] 	Loss_D: -0.1375 | Loss_G: -0.2809 | Iteration Time: 1.0056 sec
[2/200][10/32] 	Loss_D: -0.1035 | Loss_G: -0.3684 | Iteration Time: 1.3871 sec
[2/200][15/32] 	Loss_D: 0.0478 | Loss_G: -0.3379 | Iteration Time: 1.4082 sec
[2/200][20/32] 	Loss_D: -0.0658 | Loss_G: -0.3135 | Iteration Time: 0.6466 sec
[2/200][25/32] 	Loss_D: -0.0204 | Loss_G: -0.4139 | Iteration Time: 0.9207 sec

Training Progress:   2%|▏         | 3/200 [01:52<1:59:53, 36.51s/it]

[2/200][30/32] 	Loss_D: -0.1665 | Loss_G: -0.3810 | Iteration Time: 1.2556 sec
[3/200][0/32] 	Loss_D: -0.2015 | Loss_G: -0.3510 | Iteration Time: 0.9406 sec
[3/200][5/32] 	Loss_D: -0.0989 | Loss_G: -0.2984 | Iteration Time: 1.0134 sec
[3/200][10/32] 	Loss_D: -0.1801 | Loss_G: -0.4507 | Iteration Time: 1.3761 sec
[3/200][15/32] 	Loss_D: 0.2196 | Loss_G: -0.2984 | Iteration Time: 1.3361 sec
[3/200][20/32] 	Loss_D: -0.1878 | Loss_G: -0.3951 | Iteration Time: 0.6381 sec
[3/200][25/32] 	Loss_D: -0.0354 | Loss_G: -0.2700 | Iteration Time: 0.9196 sec

wgangp_trainer.plot_loss()

wgangp_trainer.plot_scores()

Minimum FID Score of 74.44414520263672 obtained at iteration of 3500.
Maximum IS Score of 2.5484938621520996.

# File directory
directory = 'wgan_weights/'
# Get file names
files = os.listdir(directory)
# Get generator files with their iteration number
generator_files = [(file, int(re.search(r'generator-(\d+)', file).group(1)))
                   for file in files if re.match(r'generator-\d+', file)]
# Sort descending by iteration number as highest number will be best weight
generator_files.sort(key=lambda x: x[1], reverse=True)
# Grab best weights file
if generator_files:
    best_weights = generator_files[0][0]
    print(f"The file with the best weights is: {best_weights}")
else:
    print("No files found.")
# Load best weights
best_wgan = torch.load(directory + best_weights)
summary(best_wgan,verbose=0,device=device)

The file with the best weights is: generator-6000.pth

=================================================================
Layer (type:depth-idx)                   Param #
=================================================================
├─Sequential: 1-1                        --
|    └─Sequential: 2-1                   --
|    |    └─ConvTranspose2d: 3-1         1,638,400
|    |    └─BatchNorm2d: 3-2             2,048
|    |    └─LeakyReLU: 3-3               --
|    └─Sequential: 2-2                   --
|    |    └─ConvTranspose2d: 3-4         8,388,608
|    |    └─BatchNorm2d: 3-5             1,024
|    |    └─LeakyReLU: 3-6               --
|    └─Sequential: 2-3                   --
|    |    └─ConvTranspose2d: 3-7         2,097,152
|    |    └─BatchNorm2d: 3-8             512
|    |    └─LeakyReLU: 3-9               --
|    └─Sequential: 2-4                   --
|    |    └─ConvTranspose2d: 3-10        524,288
|    |    └─BatchNorm2d: 3-11            256
|    |    └─LeakyReLU: 3-12              --
|    └─ConvTranspose2d: 2-5              6,144
|    └─Tanh: 2-6                         --
=================================================================
Total params: 12,658,432
Trainable params: 12,658,432
Non-trainable params: 0
=================================================================

# File directory
directory = 'wgan-gp_weights/'
# Get file names
files = os.listdir(directory)
# Get generator files with their iteration number
generator_files = [(file, int(re.search(r'generator-(\d+)', file).group(1)))
                   for file in files if re.match(r'generator-\d+', file)]
# Sort descending by iteration number as highest number will be best weight
generator_files.sort(key=lambda x: x[1], reverse=True)
# Grab best weights file
if generator_files:
    best_weights = generator_files[0][0]
    print(f"The file with the best weights is: {best_weights}")
else:
    print("No files found.")
# Load best weights
best_wgangp = torch.load(directory + best_weights)
summary(best_wgangp,verbose=0,device=device)

The file with the best weights is: generator-3500.pth

=================================================================
Layer (type:depth-idx)                   Param #
=================================================================
├─Sequential: 1-1                        --
|    └─Sequential: 2-1                   --
|    |    └─ConvTranspose2d: 3-1         819,200
|    |    └─BatchNorm2d: 3-2             1,024
|    |    └─LeakyReLU: 3-3               --
|    └─Sequential: 2-2                   --
|    |    └─ConvTranspose2d: 3-4         2,097,152
|    |    └─BatchNorm2d: 3-5             512
|    |    └─LeakyReLU: 3-6               --
|    └─Sequential: 2-3                   --
|    |    └─ConvTranspose2d: 3-7         524,288
|    |    └─BatchNorm2d: 3-8             256
|    |    └─LeakyReLU: 3-9               --
|    └─Sequential: 2-4                   --
|    |    └─ConvTranspose2d: 3-10        131,072
|    |    └─BatchNorm2d: 3-11            128
|    |    └─LeakyReLU: 3-12              --
|    └─ConvTranspose2d: 2-5              3,072
|    └─Tanh: 2-6                         --
=================================================================
Total params: 3,576,704
Trainable params: 3,576,704
Non-trainable params: 0
=================================================================

# Evaluation Metric Calculation Function
def calculate_metrics(fake, real, model_name, device='cpu'):
    # Instantiate Metrics Calculation Functions from Torchmetrics
    fid = FrechetInceptionDistance(feature=2048, normalize=True).to(device)
    inception = InceptionScore(normalize=True).to(device)

    # Move the real and fake images to CPU
    real = real.to(device)
    fake = fake.to(device)

    # FID Score calculation
    real_images_norm = GANMonitor.transform(real)
    fid.update(real_images_norm, real=True)
    fake_images_norm = GANMonitor.transform(fake)
    fid.update(fake_images_norm, real=False)
    fid_score = fid.compute().detach()

    # Inception Score calculation
    norm_imgs = GANMonitor.transform(fake.to(device))
    scaled_imgs = (norm_imgs * 255).clamp(0, 255).to(torch.uint8)
    inception.update(scaled_imgs)
    inception_score, _ = inception.compute()
    inception_score = inception_score.detach()

    # Compile the results into a DataFrame
    results_df = pd.DataFrame({
        'FID Score': [fid_score.item()],
        'Inception Score': [inception_score.item()]
    })
    results_df.rename(index={0: model_name}, inplace=True)

    # Return DataFrame
    return results_df

# Fixed noise (latent vectors)
fixed_noise = torch.randn(1000, nz, 1, 1, device=device)
# Generate images
generated_images = best_wgan(fixed_noise)
# Calculate metrics compared to test images
calculate_metrics(fake=generated_images, real=images[:1000], model_name='WGAN')

# Fixed noise (latent vectors)
fixed_noise = torch.randn(1000, nz, 1, 1, device=device)
# Generate images
generated_images = best_wgangp(fixed_noise)
# Calculate metrics compared to test images
calculate_metrics(fake=generated_images, real=images[:1000],model_name='WGAN-GP')

def generate_and_plot_single_image(best_gen, images_per_plot=64, batch_size=8, model_name="model"):
    plt.figure(figsize=(8, 8))  # Adjust figure size if needed
    image_count = 0
    while image_count < images_per_plot:
        fixed_noise = torch.randn(batch_size, nz, 1, 1, device=device)  # Ensure nz and device are defined elsewhere
        batch_images = best_gen(fixed_noise)
        size = batch_images.size(0)
        batch_images = normalize_images(batch_images)
        for i in range(size):
            if image_count >= images_per_plot:
                break
            img = batch_images[i]
            img_np = np.transpose(img.cpu().detach().numpy().clip(0, 1), (1, 2, 0))
            plt.subplot(8, 8, image_count + 1)  # Changed to 8x8 grid
            plt.imshow(img_np)
            plt.axis('off')
            image_count += 1
    plt.tight_layout()
    plt.savefig(f'./images/{model_name}_image_plot.png')
    plt.show()

generate_and_plot_single_image(best_wgan, model_name='WGAN')

generate_and_plot_single_image(best_wgangp, model_name='WGAN-GP')

ST1504 DEEP LEARNING CA2 PART C¶

IMPORT MODULES¶

PLOTTING CUSTOMIZATIONS¶

IMPORT DATA¶

CONSTANTS¶

VISUALIZE IMAGES¶

HELPER FUNCTIONS¶

WGAN¶

GENERATOR¶

CRITIC¶

OPTIMIZERS¶

TRAINING¶

LOSS PLOT¶

SCORE PLOT¶

WGAN-GP¶

GENERATOR¶

CRITIC¶

OPTIMIZERS AND HYPERPARAMETERS¶

TRAINING¶

LOSS PLOT¶

SCORE PLOT¶

EVALUATE BEST WEIGHTS¶

LOAD BEST MODELS¶

WGAN¶

WGAN-GP¶

EVALUATE QUANTITATIVE METRICS¶

EVALUATE IMAGES GENERATED¶